import matplotlib.pyplot as plt


def read_txt_file(filename):
    """
    读取TXT文件，每行一个浮点数
    返回浮点数列表
    """
    data = []
    try:
        with open(filename, 'r', encoding='utf-8') as file:
            for line in file:
                # 去除空白字符并转换为浮点数
                cleaned_line = line.strip()
                if cleaned_line and 'time' not in cleaned_line:  # 确保不是空行
                    data.append(float(cleaned_line))
        print(f"成功读取文件 {filename}，共 {len(data)} 行数据")
        return data
    except FileNotFoundError:
        print(f"错误：文件 {filename} 未找到")
        return None
    except ValueError as e:
        print(f"错误：文件 {filename} 中包含无法转换为浮点数的数据 - {e}")
        return None
    except Exception as e:
        print(f"读取文件 {filename} 时发生错误：{e}")
        return None


def read_txt_file1(filename):
    """
    读取TXT文件，每行一个浮点数
    返回浮点数列表
    """
    data = []
    try:
        with open(filename, 'r', encoding='utf-8') as file:
            for line in file:
                # 去除空白字符并转换为浮点数
                cleaned_line = line.strip()
                if cleaned_line and 'time' not in cleaned_line:  # 确保不是空行
                    data.append(float(cleaned_line)* 2)
        print(f"成功读取文件 {filename}，共 {len(data)} 行数据")
        return data
    except FileNotFoundError:
        print(f"错误：文件 {filename} 未找到")
        return None
    except ValueError as e:
        print(f"错误：文件 {filename} 中包含无法转换为浮点数的数据 - {e}")
        return None
    except Exception as e:
        print(f"读取文件 {filename} 时发生错误：{e}")
        return None


def plot_line_chart(file1_data, file2_data, file3_data, file4_data, file5_data):
    """
    绘制双折线图
    """
    # 创建图形和坐标轴
    plt.figure(figsize=(12, 6))

    # 生成X轴数据（行号，从1开始）
    x_values = range(1, len(file1_data) + 1)

    # 绘制两条折线
    plt.plot(x_values, file1_data,
             label='Zhu',
             color='black',
             marker='o',
             markersize=3,
             linewidth=1.5,
             alpha=0.8)
    plt.plot(x_values, file2_data,
             label='Huang',
             color='brown',
             marker='o',
             markersize=3,
             linewidth=1.5,
             alpha=0.8)
    plt.plot(x_values, file3_data,
             label='Le K H',
             color='green',
             marker='o',
             markersize=3,
             linewidth=1.5,
             alpha=0.8)

    plt.plot(x_values, file4_data,
             label='Ghaish',
             color='blue',
             marker='o',
             markersize=3,
             linewidth=1.5,
             alpha=0.8)

    plt.plot(x_values, file5_data,
             label='本章',
             color='red',
             marker='s',
             markersize=3,
             linewidth=1.5,
             alpha=0.8)

    # 设置图表标题和坐标轴标签
    plt.title('模型推理时间对比图', fontsize=14, fontweight='bold')
    plt.xlabel('轮次', fontsize=12)
    plt.ylabel('时间（秒）', fontsize=12)

    # 添加图例
    plt.legend(fontsize=10)

    # 添加网格线（浅灰色，提高可读性）
    plt.grid(True, alpha=0.3, linestyle='--')

    # 设置坐标轴范围
    plt.xlim(0, len(file1_data) + 1)

    # 自动调整布局
    plt.tight_layout()

    # 设置中文字体（例如黑体 SimHei）
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 或者 ['Microsoft YaHei']
    # 解决保存图像时负号'-'显示为方块的问题
    plt.rcParams['axes.unicode_minus'] = False

    # 显示图表
    # plt.show()
    plt.savefig('图17每轮次模型推理耗时统计图.jpg', dpi=600, bbox_inches='tight')


def main():
    """
    主函数：协调文件读取和图表绘制
    """
    # 定义文件名
    file1 = '../experiments_logs/time/resnet_time.txt'
    file2 = '../experiments_logs/time/cnn_lstm_attention_time.txt'
    file3 = '../experiments_logs/time/lightx3Ecg_time.txt'
    file4 = '../experiments_logs/time/ecgTrans_time.txt'
    file5 = '../experiments_logs/time/lwcet_train_time.txt'

    # 读取数据
    data1 = read_txt_file1(file1)
    data2 = read_txt_file1(file2)
    data3 = read_txt_file(file3)
    data4 = read_txt_file(file4)
    data5 = read_txt_file(file5)

    # 打印数据统计信息
    print(f"\n数据统计信息：")
    print(f"文件1 - 最小值：{min(data1):.2f}, 最大值：{max(data1):.2f}, 平均值：{sum(data1) / len(data1):.2f}")
    print(f"文件2 - 最小值：{min(data2):.2f}, 最大值：{max(data2):.2f}, 平均值：{sum(data2) / len(data2):.2f}")

    # 绘制折线图
    plot_line_chart(data1 , data2, data3, data4, data5)


# 运行程序
if __name__ == "__main__":
    main()
